
# with a data that have predictions on the same problem multiple times
# merge the prediction per problem using majority_vote (random selection for ties), preferring steps with smaller len

import sys
from com import freq, emb
from util import color

def majority_vote_0(examples):
	count = {}
	for e in examples:
		ans = e['pred'][1].strip()
		count[ans] = 1 + count.get(ans, 0)
	count = [(c,k) for k,c in count.items()]
	count.sort(reverse=True)
	ans = count[0][1]
	examples2 = [e for e in examples if ans == e['pred'][1].strip()]
	examples2.sort(key=lambda e: len(e['steps']))
	return examples2[0]

def majority_vote_1(examples, votes):
	examples2 = []
	n = len(examples) - votes + 1
	for i in range(max(n, 1)):
		examples2.append(majority_vote_0(examples[i:i+votes]))
	return examples2

def majority_vote(test_set, examples, outfile='../OUT/tmp.jsonl'):
	examples_per_problem = {}
	max_votes = 0
	for e in examples:
		pid = int(e['pid'])
		if pid not in examples_per_problem:
			examples_per_problem[pid] = []
		examples_per_problem[pid].append(e)
		max_votes = max(max_votes, len(examples_per_problem[pid]))
	print('max_votes:', max_votes)
	acc = []
	for votes in [1] + list(range(3, max_votes+1)):
		print('Votes:', votes)
		new_examples = []
		for pid, examples in examples_per_problem.items():
			new_examples.extend(majority_vote_1(examples, votes))
		emb.dump_jsonl(new_examples, outfile)
		a = freq.eval_freq(test_set, outfile)
		acc.append(a)
	print(acc)
	return acc

def load_examples(*paths):
	examples = []
	for path in paths:
		examples.extend(emb.load_jsonl(path))
	return examples

def eval_1(test_set, examples, outfile='../OUT/tmp.jsonl'):
	emb.dump_jsonl(examples, outfile)
	acc = freq.eval_freq(test_set, outfile)
	time_ = []
	prompt_tokens_ = []
	completion_tokens_ = []
	for e in examples:
		costs = e['costs']
		start_time = costs['start_time']
		end_time = costs['end_time']
		prompt_tokens = costs['prompt_tokens']
		completion_tokens = costs['completion_tokens']
		time_.append(end_time - start_time)
		prompt_tokens_.append(prompt_tokens)
		completion_tokens_.append(completion_tokens)
	time = sum(time_) / len(time_)
	prompt_tokens = sum(prompt_tokens_) / len(prompt_tokens_)
	completion_tokens = sum(completion_tokens_) / len(completion_tokens_)
	print('time:', time)
	print('prompt_tokens:', prompt_tokens)
	print('completion_tokens:', completion_tokens)
	return acc, time, prompt_tokens, completion_tokens

if __name__ == '__main__':
	test_set = sys.argv[1]
	paths = sys.argv[2:]
	examples = load_examples(*paths)
	majority_vote(test_set, examples)
	eval_1(test_set, examples)
